
#include "Function.hpp"
int Process(cv::Mat &frame, cv::Point2f box[], torch::jit::script::Module module);
float Function::pointDist(cv::Point2f a, cv::Point2f b)
{
    return sqrt((a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y));
}
void Function::findlight(cv::Mat frame, torch::jit::script::Module module)
{
    LightMatch mch;
    mch.setEnemyColor(true); //true :blue
    mch.saveImg(frame);
    mch.findPossible();
    lights = mch.possibles;
    mch.drawLights(frame);
    //own part
    int lights_num = lights.size();
    aim_deps::Light left, right;
    for (int i = 0; i < lights_num; i++)
    {
        for (int j = i + 1; j < lights_num; j++)
        {
            left = lights[i];
            right = lights[j];
            float angle_diff = abs(left.box.angle - right.box.angle);
            float length_ratio = (left.box.length - right.box.length) / (left.box.length ? right.box.length : left.box.length > right.box.length);

            float center_dist = pointDist(left.box.center, right.box.center);
            float length_ave = (left.box.length + right.box.length) / 2;
            float yDiff = abs(left.box.center.y - right.box.center.y);
            float xDiff = abs(left.box.center.x - right.box.center.x);
            float yDiff_ratio = yDiff / length_ave;
            float xDiff_ratio = xDiff / length_ave;
            float ratio = center_dist / length_ave;
            if (ratio > 6.3)
                continue;
            if (yDiff_ratio > 0.45)
                continue;
            if (xDiff_ratio < 1.35)
                continue;

            if (ratio < 0.45)
                continue;
            if (angle_diff > 5.8 || length_ratio > 0.57)
                continue;
            cv::Point2f center_real((left.box.center.x + right.box.center.x) / 2, (left.box.center.y + right.box.center.y) / 2);
            cv::Point2f center_predict01((left.box.vex[0].x + right.box.vex[1].x) / 2, (left.box.vex[0].y + right.box.vex[1].y) / 2);
            cv::Point2f center_predict02((left.box.vex[1].x + right.box.vex[0].x) / 2, (left.box.vex[1].y + right.box.vex[0].y) / 2);
            if (pointDist(center_real, center_predict01) > 4)
                continue;
            if (pointDist(center_real, center_predict02) > 4)
                continue;

            cv::Scalar color = cv::Scalar(0, 255, 0);
            cv::circle(frame, center_real, 2, cv::Scalar(200, 0, 0), 3);
            cv::Point2f left_dic = (left.box.vex[0] - left.box.vex[1]) / left.box.length;
            cv::Point2f right_dic = (right.box.vex[0] - right.box.vex[1]) / right.box.length;
            cv::Point2f box[4]; //逆时针标签
            box[0] = left.box.center + left_dic * length_ave;
            box[1] = left.box.center - left_dic * length_ave;
            box[2] = right.box.center - right_dic * length_ave;
            box[3] = right.box.center + right_dic * length_ave;
            box[0] = box[0] - cv::Point2f(8, 15);
            box[1] = box[1] + cv::Point2f(-8, 15);
            box[2] = box[2] + cv::Point2f(8, 15);
            box[3] = box[3] - cv::Point2f(-8, 15);
            cv::line(frame, box[0], box[1], color, 3, cv::LINE_8);
            cv::line(frame, box[2], box[3], color, 3, cv::LINE_8);
            cv::line(frame, box[1], box[2], color, 3, cv::LINE_8);
            cv::line(frame, box[3], box[0], color, 3, cv::LINE_8);
            int num = Process(frame, box, module);
            cv::Point2f puttxt = right.box.vex[0] + cv::Point2f(10, -10);
            putText(frame, std::to_string(num), puttxt,
                    cv::FONT_HERSHEY_PLAIN, 2.0, cv::Scalar(0, 255, 0), 2); //参数一目标图像，参数二文本，参数三文本位置，参数四字体类型， 参数五字体大小，参数六字体颜色，参数七文本厚度
            return;
        }
    }
}
int Process(cv::Mat &frame, cv::Point2f box[], torch::jit::script::Module module)
{

    cv::Scalar color = cv::Scalar(0, 0, 0);
    cv::Rect roi(box[0].x, box[0].y, box[3].x - box[0].x, box[1].y - box[0].y);
    cv::Mat image_roi = frame(roi); //裁减
    cv::Mat img_gray;
    cv::resize(image_roi, img_gray, cv::Size(100, 100));
    cv::cvtColor(img_gray, img_gray, cv::COLOR_RGB2GRAY);
    cv::Mat img;
    cv::threshold(img_gray, img, 8, 255, 1);

    //load .pt
    cv::cvtColor(img, img, cv::COLOR_GRAY2RGB);

    std::vector<int64_t> sizes = {1, img.rows, img.cols, 3};
    torch::Tensor img_tensor = torch::from_blob(img.data, torch::IntList(sizes), torch::kByte);
    // img_tensor = torch::rand(0,3,100,100);
    img_tensor.unsqueeze(0);
    img_tensor = img_tensor.permute({0, 3, 1, 2}).toType(torch::kFloat32);
    img_tensor = img_tensor.div(255); //将上述数字到0——1
    img_tensor[0][0] = img_tensor[0][0].sub_(0.485).div_(0.229);
    img_tensor[0][1] = img_tensor[0][1].sub_(0.456).div_(0.224);
    img_tensor[0][2] = img_tensor[0][2].sub_(0.406).div_(0.225);
    torch::Tensor output = module.forward({img_tensor}).toTensor();
    output = torch::softmax(output, 1);

    int value = torch::argmax(output).item<int>();
    return value + 1;
}